#include <vector>
#include <iomanip> 
#include <iostream>
#include <chrono>
#include "matrix_operations_cpu.h"
#include "matrix_operations_gpu.cu"
#include "simple_fnn.h"
#include "relu.h"
#include "sigmoid.h"
#include "loss_function.h"

int main(int argc, char *argv[]) {
    if (argc < 2) {
        std::cout << "unknow mode" << std::endl;
        return 1;
    }
    std::shared_ptr<MatrixOperations> matrixOperations; //矩阵操作类
    int mode = std::atoi(argv[1]);
    if(mode == 0){
        matrixOperations = std::make_shared<MatrixOperationsGPU>();
    }else if(mode == 1){
        matrixOperations = std::make_shared<MatrixOperationsCPU>();
    }else{
        std::cout << "unknow mode: " << mode << std::endl;
        return 1;
    }
    
    auto start = std::chrono::high_resolution_clock::now();
    auto lossFunction = std::make_shared<MeanSquaredErrorLoss>(); //损失函数类
    auto activationFunction = std::make_shared<Sigmoid>();
    // 构建神经网络的结构
    size_t inputSize = 2; // 输入层大小
    size_t hiddenLayerSize = 4; // 隐藏层大小
    size_t numHiddenLayers = 4; //隐藏层数量
    size_t outputSize = 1; // 输出层大小
    float learningRate = 0.001;

    FNN myFNN(inputSize, hiddenLayerSize, numHiddenLayers, outputSize, learningRate, activationFunction, matrixOperations, lossFunction);

    // 构建简单的训练数据集  XOR 问题的期望输出
    // auto result = generateXORData(1000);
    // std::vector<std::vector<float>> inputs = result.first;
    // std::vector<std::vector<float>> expectedOutputs = result.second;
    std::vector<std::vector<float>> inputs = {{0, 0}, {0, 1}, {1, 0}, {1, 1}};
    std::vector<std::vector<float>> expectedOutputs = {{0}, {1}, {1}, {0}};

    // 训练过程
    int epochs = 100; // 训练轮数
    for (int epoch = 0; epoch < epochs; epoch++) {
        float totalLoss = 0.0;
        for (size_t j = 0; j < inputs.size(); j++) {
            myFNN.train({inputs[j]}, {expectedOutputs[j]});
            auto predictions = myFNN.predict({inputs[j]});
            totalLoss += myFNN.lossFunction->loss(predictions, {expectedOutputs[j]});
        }
        totalLoss /= inputs.size();
        std::cout << "Epoch " << (epoch + 1) << "/" << epochs << ", Loss: " << totalLoss << std::endl;
    }

    // 测试神经网络的表现
    for (size_t i = 0; i < inputs.size(); i++) {
        if(i > 10) break;
        auto output = myFNN.predict({inputs[i]});
        // Apply threshold
        float binaryOutput = output[0][0] >= 0.5 ? 1 : 0;
        std::cout << "Input: " << inputs[i][0] << ", " << inputs[i][1]
                << " - Expected Output: " << expectedOutputs[i][0]
                << " - Binary Output: " << binaryOutput
                << " - Predicted Output: " << output[0][0] << std::endl;
    }

    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(stop - start);
    std::cout << "Time taken by function: " << duration.count() << " milliseconds" << std::endl;
    fflush(stdout);
    return 0;
}
